import json

import fitz

from app.utils.cv_tools import group_text_lines_by_proximity, detect_overlapping_text_lines
from app.utils.file_tools import extract_filename_without_extension
from app.utils.pymupdf_tools import pdf_get_content_region, get_pages_from_range
from app.utils.timing import timing_decorator
from tests.base_test import base_test_case

logger = base_test_case.get_logger(__name__)
TEST_DATA_DIR = base_test_case.test_data_dir
OUTPUT_DATA_DIR = base_test_case.output_data_dir


@timing_decorator()
def load_pdf(input_pdf=None, page_range=None):
    zoom = 1.0
    filename_without_ext = f"{extract_filename_without_extension(input_pdf)}_v3"

    doc = fitz.open(input_pdf)
    try:
        pages = get_pages_from_range(input_pdf, page_range=page_range)
        ocr_results = []
        for page in pages:
            page_num = page.number
            # 获取页内容区域
            clip_rect = pdf_get_content_region(page, zoom_factor=zoom, is_show_log=False)
            block_dict = page.get_text("dict", clip=clip_rect, sort=True)

            for block in block_dict.get("blocks", []):
                for line in block.get("lines", []):
                    for span in line.get("spans", []):
                        text = span["text"]
                        bbox = fitz.Rect(span["bbox"])
                        ocr_results.append({
                            "text": text,
                            "bbox": list(bbox.irect),
                            "height": bbox.irect.height,
                            "width": bbox.irect.width,
                            "size": span["size"],
                            "flags": span["flags"],
                            "bidi": span["bidi"],
                            "char_flags": span["char_flags"],
                            "font": span["font"],
                            "color": span["color"],
                            "alpha": span["alpha"],
                            "ascender": span["ascender"],
                            "descender": span["descender"],
                        })
            output_path = OUTPUT_DATA_DIR / f"{filename_without_ext}_{page_num:02d}.json"
            json_str = json.dumps(ocr_results, indent=4, ensure_ascii=False)
            if output_path:
                with open(output_path, "w", encoding="utf-8") as file:
                    file.write(json_str)

        return ocr_results, filename_without_ext, page_num
    finally:
        doc.close()


@timing_decorator()
def analyze_text_structure(ocr_results,
                           header_margin_ratio=0.05,
                           footer_margin_ratio=0.05,
                           proximity_vertical_threshold=20,
                           overlap_threshold=0.3,
                           thresholds=None):
    """
    使用提供的函数分析文本结构，识别段落、标题等
    
    Args:
        ocr_results: OCR识别结果列表
        header_margin_ratio: 页眉区域比例（相对于页面高度），默认0.05（5%）
        footer_margin_ratio: 页脚区域比例（相对于页面高度），默认0.05（5%）
        proximity_vertical_threshold: 文本行 proximity 分组的垂直距离阈值，默认20
        overlap_threshold: 文本行重叠检测的阈值，默认0.3
        thresholds: 各类文本元素的阈值配置字典，默认为None，使用内置默认值
        
    Returns:
        dict: 包含分析结果的字典
    """
    # 默认阈值配置，参考PP-StructureV3的配置方式
    default_thresholds = {
        # 按元素类型分类的阈值
        'paragraph_title': 0.3,  # 段落标题
        'image': 0.5,  # 图片
        'text': 0.4,  # 正文文本
        'number': 0.5,  # 数字
        'abstract': 0.5,  # 摘要
        'content': 0.5,  # 内容
        'figure_table_chart_title': 0.5,  # 图表标题
        'formula': 0.3,  # 公式
        'table': 0.5,  # 表格
        'reference': 0.5,  # 参考文献
        'doc_title': 0.5,  # 文档标题
        'footnote': 0.5,  # 脚注
        'header': 0.5,  # 页眉
        'algorithm': 0.5,  # 算法
        'footer': 0.5,  # 页脚
        'seal': 0.45,  # 印章
        'chart': 0.5,  # 图表
        'formula_number': 0.5,  # 公式编号
        'aside_text': 0.5,  # 旁注文本
        'reference_content': 0.5,  # 参考文献内容

        # 标题检测相关阈值
        'title_size_ratio': 1.5,  # 标题相对于平均字号的倍数阈值
        'title_min_size': 16,  # 标题的最小绝对大小
        'large_title_position_ratio': 0.2,  # 大标题在页面中的位置比例阈值
        'section_title_position_ratio': 0.5,  # 小节标题在页面中的位置比例阈值
        'section_title_size_ratio': 1.4,  # 小节标题相对于平均字号的倍数阈值
        'main_title_size_ratio': 1.3,  # 主标题相对于平均字号的倍数阈值
    }

    # 如果提供了自定义阈值，则合并到默认阈值中
    if thresholds:
        default_thresholds.update(thresholds)

    # 提取阈值变量以便使用
    title_size_ratio = default_thresholds['title_size_ratio']
    title_min_size = default_thresholds['title_min_size']
    large_title_position_ratio = default_thresholds['large_title_position_ratio']
    section_title_position_ratio = default_thresholds['section_title_position_ratio']
    section_title_size_ratio = default_thresholds['section_title_size_ratio']
    main_title_size_ratio = default_thresholds['main_title_size_ratio']

    # 提取特定元素类型的阈值
    paragraph_title_threshold = default_thresholds['paragraph_title']
    formula_threshold = default_thresholds['formula']
    footnote_threshold = default_thresholds['footnote']
    aside_text_threshold = default_thresholds['aside_text']

    # 将数据转换为所需的格式
    # ocr_results 格式: [[line1, line2, ...]]
    # line 格式: [box, [text, confidence], ...]
    formatted_results = [[
        [
            [  # box - 四个点的坐标
                [item['bbox'][0], item['bbox'][1]],  # 左上角
                [item['bbox'][2], item['bbox'][1]],  # 右上角
                [item['bbox'][2], item['bbox'][3]],  # 右下角
                [item['bbox'][0], item['bbox'][3]]  # 左下角
            ],
            [item['text'], 1.0]  # [text, confidence]
        ] for item in ocr_results
    ]]

    # 使用group_text_lines_by_proximity按垂直距离对文本进行分组
    proximity_groups = group_text_lines_by_proximity(formatted_results, vertical_threshold=proximity_vertical_threshold)

    # 检测重叠的文本行
    overlapping_pairs = detect_overlapping_text_lines(formatted_results, overlap_threshold=overlap_threshold)

    # 计算高度和位置的统计信息，用于判断标题
    heights = [item['height'] for item in ocr_results if item['text'].strip()]
    sizes = [item.get('size', item['height']) for item in ocr_results if item['text'].strip()]

    if heights:
        avg_height = sum(heights) / len(heights)
        max_height = max(heights)
        min_height = min(heights)
    else:
        avg_height = max_height = min_height = 0

    if sizes:
        avg_size = sum(sizes) / len(sizes)
        max_size = max(sizes)
        min_size = min(sizes)
    else:
        avg_size = max_size = min_size = 0

    # 按y坐标对文本进行排序，以便分析页面布局
    sorted_items = sorted([item for item in ocr_results if item['text'].strip()],
                          key=lambda x: x['bbox'][1])

    # 分析文本结构
    paragraphs = []
    titles = []
    headers = []
    footers = []
    formulas = []
    footnotes = []
    aside_texts = []
    content_lines = []

    # 确定页面边界
    if sorted_items:
        page_top = min(item['bbox'][1] for item in sorted_items)
        page_bottom = max(item['bbox'][3] for item in sorted_items)
        page_height = page_bottom - page_top

        # 定义页眉页脚区域
        header_zone_bottom = page_top + page_height * header_margin_ratio
        footer_zone_top = page_bottom - page_height * footer_margin_ratio

        # 定义旁注区域（页面左右边缘区域）
        min_x = min(item['bbox'][0] for item in sorted_items)
        max_x = max(item['bbox'][2] for item in sorted_items)
        page_width = max_x - min_x
        aside_zone_width = page_width * 0.1  # 页面左右各10%为旁注区域
        aside_left_zone = min_x + aside_zone_width
        aside_right_zone = max_x - aside_zone_width
    else:
        page_top = page_bottom = header_zone_bottom = footer_zone_top = 0
        aside_left_zone = aside_right_zone = 0

    for item in sorted_items:
        text = item['text'].strip()
        height = item['height']
        size = item.get('size', height)
        width = item['width']
        y_position_top = item['bbox'][1]  # 文本的y坐标（顶部）
        y_position_bottom = item['bbox'][3]  # 文本的y坐标（底部）
        x_position_left = item['bbox'][0]  # 文本的x坐标（左侧）
        x_position_right = item['bbox'][2]  # 文本的x坐标（右侧）

        if not text or text.isspace():
            continue

        # 判断是否为页眉或页脚
        is_header_footer = False

        # 页眉：位于页面顶部区域的文本
        if y_position_bottom <= header_zone_bottom:
            headers.append(item)
            is_header_footer = True
            continue

        # 页脚：位于页面底部区域的文本
        if y_position_top >= footer_zone_top:
            footers.append(item)
            is_header_footer = True
            continue

        # 判断是否为旁注文本（位于页面左右边缘区域）
        if x_position_right <= aside_left_zone or x_position_left >= aside_right_zone:
            # 使用旁注文本阈值进行判断
            if len(text) > 0 and len(text) < 50:  # 简单的长度判断
                aside_texts.append(item)
                continue

        # 判断是否为脚注（页面底部但在页脚区域之上，且字号较小）
        if (y_position_top >= page_bottom - page_height * 0.3 and  # 在页面底部30%区域
                y_position_top < footer_zone_top and  # 但在页脚区域之上
                size < avg_size * 0.8):  # 字号相对较小
            footnotes.append(item)
            continue

        # 判断是否为公式（包含特殊数学符号）
        math_symbols = ['∫', '∑', '∏', '√', '∞', '≠', '≤', '≥', '±', '∓', '×', '÷', '∂', '∆', '∇']
        if any(symbol in text for symbol in math_symbols):
            # 使用公式阈值进行判断
            formulas.append(item)
            continue

        # 判断是否为标题 - 主要基于视觉特征和字体信息
        is_title = False

        # 基于字号判断（显著大于平均字号的可能是标题）
        if size >= avg_size * title_size_ratio:
            is_title = True

        # 基于绝对字号判断（特别大的文字）
        if size >= title_min_size:
            is_title = True

        # 如果有字体信息，粗体或特殊字体可能是标题
        font = item.get('font', '')
        if 'bold' in font.lower() or 'title' in font.lower() or 'heading' in font.lower():
            is_title = True

        # 页面顶部的大号文字很可能是章节标题
        if y_position_top < page_top + page_height * large_title_position_ratio and size > avg_size * main_title_size_ratio:
            is_title = True

        # 页面中的较大文字可能是小节标题
        if (page_top + page_height * large_title_position_ratio <= y_position_top <
                page_top + page_height * section_title_position_ratio and
                size > avg_size * section_title_size_ratio):
            is_title = True

        if is_title:
            titles.append(item)
        else:
            content_lines.append(item)

    # 使用 proximity_groups 来组织段落
    for group in proximity_groups:
        if group:  # 确保组不为空
            # 按垂直位置排序组内的文本行
            group.sort(key=lambda x: x['center'][1])

            # 将组内文本连接成段落文本
            paragraph_text = ''.join([line['text'] for line in group])

            # 获取段落的边界框
            min_top = min(line['top'] for line in group)
            max_bottom = max(line['bottom'] for line in group)

            # 计算段落的近似左边界和右边界
            min_left = min(min(point[0] for point in line['box']) for line in group)
            max_right = max(max(point[0] for point in line['box']) for line in group)

            paragraphs.append({
                'text': paragraph_text,
                'lines': group,
                'bbox': [min_left, min_top, max_right, max_bottom],
                'line_count': len(group)
            })

    return {
        'proximity_groups': proximity_groups,
        'overlapping_pairs': overlapping_pairs,
        'titles': titles,
        'paragraphs': paragraphs,
        'headers': headers,
        'footers': footers,
        'formulas': formulas,
        'footnotes': footnotes,
        'aside_texts': aside_texts,
        'content_lines': content_lines,
        'stats': {
            'avg_height': avg_height,
            'max_height': max_height,
            'min_height': min_height,
            'avg_size': avg_size,
            'max_size': max_size,
            'min_size': min_size,
            'page_top': page_top,
            'page_bottom': page_bottom,
            'page_height': page_bottom - page_top if page_bottom > page_top else 0,
            'header_margin_ratio': header_margin_ratio,
            'footer_margin_ratio': footer_margin_ratio,
            'proximity_vertical_threshold': proximity_vertical_threshold,
            'overlap_threshold': overlap_threshold,
            'thresholds': default_thresholds
        }
    }


def print_analyze_text_structure(analysis_result):
    # 输出分析结果
    print("\n=== 文本结构分析结果 ===")
    print(f"检测到标题数量: {len(analysis_result['titles'])}")
    for title in analysis_result['titles']:
        size = title.get('size', title['height'])
        font = title.get('font', 'Unknown')
        print(
            f"  标题: '{title['text']}' (大小: {size:.1f}, 字体: {font}, 位置: ({title['bbox'][0]}, {title['bbox'][1]}))")

    print(f"\n检测到页眉数量: {len(analysis_result['headers'])}")
    for header in analysis_result['headers']:
        print(f"  页眉: '{header['text']}'")

    print(f"\n检测到页脚数量: {len(analysis_result['footers'])}")
    for footer in analysis_result['footers']:
        print(f"  页脚: '{footer['text']}'")

    print(f"\n检测到公式数量: {len(analysis_result['formulas'])}")
    for formula in analysis_result['formulas']:
        print(f"  公式: '{formula['text']}'")

    print(f"\n检测到脚注数量: {len(analysis_result['footnotes'])}")
    for footnote in analysis_result['footnotes']:
        print(f"  脚注: '{footnote['text']}'")

    print(f"\n检测到旁注数量: {len(analysis_result['aside_texts'])}")
    for aside in analysis_result['aside_texts']:
        print(f"  旁注: '{aside['text']}'")

    print(f"\n检测到段落数量: {len(analysis_result['paragraphs'])}")
    for i, paragraph in enumerate(analysis_result['paragraphs']):
        print(f"  段落 {i + 1}:")
        print(f"    内容: '{paragraph['text']}'")
        print(f"    行数: {paragraph['line_count']}")
        print(f"    位置: {paragraph['bbox']}")

    print(f"\n检测到相近文本组数量: {len(analysis_result['proximity_groups'])}")
    print(f"检测到重叠文本对数量: {len(analysis_result['overlapping_pairs'])}")

    # 打印统计信息
    stats = analysis_result['stats']
    print(f"\n文档统计信息:")
    print(f"  平均文字高度: {stats['avg_height']:.2f}")
    print(f"  平均文字大小: {stats['avg_size']:.2f}")
    print(f"  最大文字大小: {stats['max_size']:.2f}")
    print(f"  最小文字大小: {stats['min_size']:.2f}")


if __name__ == "__main__":
    # input_pdf = TEST_DATA_DIR / "25-注会-轻1-财务成本管理[上册](第3章).pdf"
    # input_pdf = TEST_DATA_DIR / "1715340454803.pdf"
    # input_pdf = TEST_DATA_DIR / "1715306099910.pdf"
    # input_pdf = TEST_DATA_DIR / "1711605374231.pdf"
    input_pdf = TEST_DATA_DIR / "1715339805571.pdf"
    page_range = (1, 1)
    ocr_results, filename_without_ext, page_num = load_pdf(input_pdf, page_range)
    # 使用group_text_lines_by_proximity,detect_overlapping_text_lines,calculate_rectangle_center 函数分析出文本的段落章节，标题等
    # 自定义所有阈值（类似PP-StructureV3的配置方式）
    all_thresholds = {
        'title_size_ratio': 1.5,
        'title_min_size': 16,
        'large_title_position_ratio': 0.2,
        'section_title_position_ratio': 0.5,
        'section_title_size_ratio': 1.4,
        'main_title_size_ratio': 1.3,
    }
    analysis_result = analyze_text_structure(ocr_results,
                                             header_margin_ratio=0.0,
                                             footer_margin_ratio=0.0,
                                             thresholds=all_thresholds)

    print_analyze_text_structure(analysis_result)

    # 保存分析结果
    analysis_output_path = OUTPUT_DATA_DIR / f"{filename_without_ext}_{page_num:02d}_analysis.json"
    with open(analysis_output_path, "w", encoding="utf-8") as file:
        json.dump(analysis_result, file, indent=4, ensure_ascii=False,
                  default=lambda o: '<not serializable>' if isinstance(o, (fitz.Rect,)) else o)
